﻿#ifndef __COLOR_TRAINER_H__
#define __COLOR_TRAINER_H__

#include "SoftMaxTrainer.h"

class ColorTrainer : public SoftMaxTrainer
{
public:
    static ColorTrainer *GetInstance()
    {
        static ColorTrainer s_instance;
        static bool s_loaded = false;

        if (!s_loaded)
        {
#ifdef __EMSCRIPTEN__
            std::string colors_model_path = "/models/colors.model";
#else
            std::string colors_model_path = "D:/workspace/cpp/darknet/web/asset/models/colors.model";
#endif
            if (!s_instance.LoadModelFromFile(colors_model_path))
            {
                std::cout << "invalid colors model path: " << colors_model_path << std::endl;
            }

            s_loaded = true;
        }

        return &s_instance;
    }

public:
    ColorTrainer()
    {
        m_max_epochs      = 10000;
        m_batch           = 100;
        m_auto_stop_train = false;
        m_models_path     = "D:/workspace/cpp/darknet/data/colors/models";

        CreateGetClassifyIndex("red");
        CreateGetClassifyIndex("blank");
        CreateGetClassifyIndex("white");
        CreateGetClassifyIndex("other");
    }

    void LoadNet()
    {
        auto net     = GetNet();
        auto options = std::make_shared<std::vector<std::shared_ptr<Options>>>();
        auto option  = std::make_shared<Options>("net");
        option->SetOption("width", 5);
        option->SetOption("height", 5);
        option->SetOption("channels", 3);
        options->push_back(option);

        option = std::make_shared<Options>("convolutional");
        option->SetOption("size", 3);
        option->SetOption("stride", 1);
        option->SetOption("padding", 0);
        option->SetOption("filters", 16);
        options->push_back(option);

        option = std::make_shared<Options>("connected");
        option->SetOption("output", 32);
        option->SetOption("activation", "leaky");
        options->push_back(option);

        option = std::make_shared<Options>("connected");
        option->SetOption("output", 64);
        option->SetOption("activation", "leaky");
        options->push_back(option);

        option = std::make_shared<Options>("connected");
        option->SetOption("output", 32);
        option->SetOption("activation", "leaky");
        options->push_back(option);

        option = std::make_shared<Options>("connected");
        option->SetOption("output", GetClassifySize());
        option->SetOption("activation", "linear");
        options->push_back(option);

        option = std::make_shared<Options>("softmax");
        options->push_back(option);

        net->Load(options);
    }

    std::string PredictByPath(const std::string &path)
    {
        auto img = std::make_shared<DataImage>();
        if (!img->Load(path)) return "";
        img = img->ScaleFill(GetNet()->GetWidth(), GetNet()->GetHeight());
        return GetIndexClassify(Predict(img->GetData()));
    }

    std::string PredictColor(unsigned char r, unsigned char g, unsigned char b)
    {
        auto width  = GetNet()->GetWidth();
        auto height = GetNet()->GetHeight();
        auto img    = std::make_shared<DataImage>(width, height, 3);
        for (int x = 0; x < width; x++)
        {
            for (int y = 0; y < height; y++)
            {
                img->SetPixel(x, y, 0, r);
                img->SetPixel(x, y, 1, g);
                img->SetPixel(x, y, 2, b);
            }
        }

        return GetIndexClassify(Predict(img->GetData()));
    }

    int Predict(float *input)
    {
        auto net = GetNet();
        net->SetBatch(1);
        net->SetInDatas(input);
        net->SetInTruths(nullptr);
        net->Forward();
        auto layer     = net->GetBackLayer();
        auto results   = layer->GetPredicts();
        auto size      = results->size();
        auto max_index = 0;
        auto max_value = results->at(0);
        for (int i = 1; i < size; i++)
        {
            if (results->at(i) > max_value)
            {
                max_index = i;
                max_value = results->at(i);
            }
        }
        // std::cout << "index: " << max_index << ", value: " << max_value << "label: " << m_index_classify_map[max_index] << std::endl;
        return max_index;
    }

    void Train()
    {
        m_epochs             = 0;
        auto net             = GetNet();
        auto os              = GetTrainOS();
        auto max_epochs      = GetMaxEpcohs();
        auto auto_stop_train = IsAutoStopTrain();
        auto batch           = GetBatch();
        int costs_size       = m_costs.size();
        float last_avg_cost  = FLT_MAX;

        net->SetBatch(batch);
        PrintTrainerInfo(os);

        while (m_epochs < max_epochs)
        {
            auto start_ts = get_timestamp();

            // 获取下一批数据
            auto input = GetNextBatchTrainInput(batch);
            if (input == nullptr) return;
            // 设置输入输出流
            net->SetInOSS(input->GetInOSS());

            // 设置输入数据
            net->SetInDatas(input->GetDatas());
            net->SetInTruths(input->GetTruths());

            // 前向网络
            net->Forward();
            // 计算损失
            auto cost = net->Loss();
            // 反向网络
            net->Backward();

            // 更新网络
            net->Update(GetCurrentLearningRate(), GetMomentum(), GetDecay(), GetEpochBatchSize());

            auto end_ts                    = get_timestamp();
            m_costs[m_epochs % costs_size] = cost;
            m_epochs++;

            *os << "epochs: " << m_epochs << " use times:" << (end_ts - start_ts) << " ms loss: " << cost << std::endl;
            if (os != &std::cout) std::cout << "epochs: " << m_epochs << " use times:" << (end_ts - start_ts) << " ms loss: " << cost << std::endl;

            if (auto_stop_train && (m_epochs % costs_size) == 0)
            {
                auto cur_avg_cost   = GetAverageCost();
                auto delta_avg_cost = std::abs(last_avg_cost - cur_avg_cost);
                *os << "per 10 epochs: cur_avg_cost = " << cur_avg_cost << " last_avg_cost = " << last_avg_cost << " delta_avg_cost = " << delta_avg_cost << std::endl;
                if (delta_avg_cost < 0.0001)
                {
                    *os << "auto stop train!!!" << std::endl;
                    break;
                }
                last_avg_cost = cur_avg_cost;
            }
        }

        auto models_directory = std::filesystem::path(GetModelsPath());
        if (!models_directory.empty()) net->Save((models_directory / m_default_model_filename).string());
    }

    std::shared_ptr<DataSet::BatchInput> GetNextBatchTrainInput(int batch)
    {
        auto net                = GetNet();
        auto datas_size         = net->GetWidth() * net->GetHeight() * net->GetChannels();
        auto truths_size        = net->GetBackLayer()->GetOutSize();
        auto batch_input        = std::make_shared<DataSet::BatchInput>(batch * datas_size, batch * truths_size);
        auto batch_input_datas  = batch_input->GetDatas();
        auto batch_input_truths = batch_input->GetTruths();

        auto size   = m_train_truths.size();
        auto datas  = m_train_datas.data();
        auto truths = m_train_truths.data();
        for (int i = 0; i < batch; i++)
        {
            int index = random_int(0, size - 1);
            for (int j = 0; j < datas_size; j++)
            {
                batch_input_datas[i * datas_size + j] = datas[index * datas_size + j];
            }
            int truth                                   = truths[index];
            batch_input_truths[i * truths_size + truth] = 1.f;
        }

        return batch_input;
    }

    void LoadTrainData()
    {
        std::string train_data_path = "D:/workspace/cpp/darknet/data/colors/train.txt";
        std::ifstream ifs(train_data_path);

        std::string line;
        while (ifs.good() && std::getline(ifs, line))
        {
            auto pos = line.find(" ");
            if (pos == std::string::npos) continue;
            auto label = trim(line.substr(0, pos));
            auto path  = trim(line.substr(pos + 1));
            if (label.empty() || path.empty()) continue;
            auto label_index = CreateGetClassifyIndex(label);
            auto img         = std::make_shared<DataImage>();
            if (!img->Load(path)) continue;

            auto img_width    = img->GetWidth();
            auto img_height   = img->GetHeight();
            auto img_channels = img->GetChannel();
            auto img_datas    = img->GetData();
            auto net_width    = m_net->GetWidth();
            auto net_height   = m_net->GetHeight();
            for (int j = 0; j < img_height; j += net_height)
            {
                for (int i = 0; i < img_width; i += net_width)
                {
                    auto sub_img = img->GetSubImage(i, j, net_width, net_height);
                    m_train_datas.insert(m_train_datas.end(), sub_img->GetData(), sub_img->GetData() + sub_img->GetSize());
                    m_train_truths.push_back(label_index);
                }
            }
        }
    }

    std::shared_ptr<DataImage> GetColorImage(std::shared_ptr<DataImage> img, int color_classify_index, int max_distance = 200)
    {
        std::vector<int> xs;
        std::vector<int> ys;
        auto img_width    = img->GetWidth();
        auto img_height   = img->GetHeight();
        auto img_channels = img->GetChannel();
        auto img_datas    = img->GetData();
        auto net          = GetNet();
        auto net_width    = net->GetWidth();
        auto net_height   = net->GetHeight();

        for (int y = 0; y < img_height; y += net_height)
        {
            for (int x = 0; x < img_width; x += net_width)
            {
                auto sub_img = img->GetSubImage(x, y, net_width, net_height);
                auto index   = Predict(sub_img->GetData());
                if (color_classify_index == index)
                {
                    xs.push_back(x);
                    ys.push_back(y);
                }
            }
        }

        int size        = xs.size();
        auto avg_x      = 0;
        auto avg_y      = 0;
        auto min_x      = img_width;
        auto min_y      = img_height;
        auto max_x      = 0;
        auto max_y      = 0;
        auto valid_size = 0;
        for (int i = 0; i < size; i++)
        {
            auto x     = xs[i];
            auto y     = ys[i];
            auto count = 0;

            for (int j = 0; j < size; j++)
            {
                if (i == j) continue;
                if (std::abs(x - xs[j]) < max_distance && std::abs(y - ys[j]) < max_distance) count++;
            }

            if (count > (size * 2 / 3))
            {
                avg_x += x;
                avg_y += y;
                min_x = min_x < x ? min_x : x;
                min_y = min_y < y ? min_y : y;
                max_x = max_x < x ? x : max_x;
                max_y = max_y < y ? y : max_y;
                // std::cout << "x = " << x << " y = " << y << std::endl;
                valid_size++;
            }
        }
        if (min_x > max_x || min_y > max_y) return nullptr;

        avg_x /= valid_size;
        avg_y /= valid_size;

        int half_width  = (max_x - min_x) / 2;
        int half_height = (max_y - min_y) / 2;
        int step_x      = 0;
        int step_y      = 0;
        int rect[4]     = {0};
        do {
            if (rect[0] > min_x || rect[1] < max_x) step_x += 5;
            if (rect[2] > min_y || rect[3] < max_y) step_y += 5;
            rect[0] = avg_x - half_width - step_x;
            rect[1] = avg_x + half_width + step_x;
            rect[2] = avg_y - half_height - step_y;
            rect[3] = avg_y + half_height + step_y;
        } while (rect[0] > min_x || rect[1] < max_x || rect[2] > min_y || rect[3] < max_y);
        // std::cout << "min_x = " << min_x << ", min_y = " << min_y << ", max_x = " << max_x << ", max_y = " << max_y << ", avg_x = " << avg_x << ", avg_y = " << avg_y << ", size = " << valid_size << std::endl;

        auto subimg = img->GetSubImage(rect[0], rect[2], rect[1] - rect[0], rect[3] - rect[2]);
        // for (int i = 0; i < size; i++)
        // {
        //     int x = xs[i] - rect[0];
        //     int y = ys[i] - rect[2];
        //     subimg->FillRectangle(x, y, net_width, net_height, 0.f, 0.f, 0.f, 0.f);
        // }
        return subimg;
    }

protected:
    std::string trim(std::string str)
    {
        // clang-format off
        // 移除开头的空白字符
        str.erase(str.begin(), std::find_if(str.begin(), str.end(), [](unsigned char ch) { return !std::isspace(ch); }));

        // 移除结尾的空白字符
        str.erase(std::find_if(str.rbegin(), str.rend(), [](unsigned char ch) { return !std::isspace(ch); }).base(),str.end());
        // clang-format on
        return str;
    }

protected:
    std::vector<float> m_train_datas;
    std::vector<int> m_train_truths;
};

#endif